import argparse

import torch
import torch.backends.cudnn as cudnn
import numpy as np
import PIL.Image as pil_image
from PIL import Image

from models import ResAttU_Net,UNet,ResAttU_Net_noBN
from utils import convert_rgb_to_ycbcr, convert_ycbcr_to_rgb, calc_psnr

weights_file = r'权重保存路径\best.pth'
image_file  = r'待处理lr图片'  # 在这里输入lr的照片
# pic = Image.open(image_file1)
# image_file = pic.resize((512, 512))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # parser.add_argument('--weights-file', type=str, default=weights_file, required=True) # lcj
    # parser.add_argument('--image-file', type=str,  default= test_image , required=True) # lcj
    parser.add_argument('--weights-file', type=str, default=weights_file)
    parser.add_argument('--image-file', type=str, default=image_file)
    # parser.add_argument('--scale', type=int, default=2)
    args = parser.parse_args()

    cudnn.benchmark = True
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    ### ****** 训练时使用GPU，使用了torch.nn.DataParallel()，而此时预测没有使用GPU，更改如下：*****
    model_DP = torch.nn.DataParallel(ResAttU_Net_noBN())

    model = model_DP.to(device)
    # model = UNet().to(device)  -- lcj_5.17

    ######

    state_dict = model.state_dict()
    for n, p in torch.load(args.weights_file, map_location=lambda storage, loc: storage).items():
        if n in state_dict.keys():
            state_dict[n].copy_(p)
        else:
            raise KeyError(n)

    model.eval()

    image = pil_image.open(args.image_file).convert('RGB')


    # image_width = (image.width // args.scale) * args.scale   # lcj 3.23更改
    # image_height = (image.height // args.scale) * args.scale # lcj 3.23更改
    # image = image.resize((image_width, image_height), resample=pil_image.BICUBIC)
    # image = image.resize((image.width // args.scale, image.height // args.scale), resample=pil_image.BICUBIC)
    # image = image.resize((image.width * args.scale, image.height * args.scale), resample=pil_image.BICUBIC)
    # image.save(args.image_file.replace('.', '_bicubic_x{}.'.format(args.scale)))

    image = np.array(image).astype(np.float32)
    ycbcr = convert_rgb_to_ycbcr(image)

    y = ycbcr[..., 0]
    y /= 255.
    y = torch.from_numpy(y).to(device)
    y = y.unsqueeze(0).unsqueeze(0)

    with torch.no_grad():
        preds = model(y).clamp(0.0, 1.0)

    psnr = calc_psnr(y, preds)
    print('PSNR: {:.2f}'.format(psnr))

    preds = preds.mul(255.0).cpu().numpy().squeeze(0).squeeze(0)

    output = np.array([preds, ycbcr[..., 1], ycbcr[..., 2]]).transpose([1, 2, 0])
    output = np.clip(convert_ycbcr_to_rgb(output), 0.0, 255.0).astype(np.uint8)
    output = pil_image.fromarray(output)
    # output.save(args.image_file.replace('.', '_srcnn_x{}_PSNR{:.2f}.'.format(args.scale, psnr)))
    output.save(args.image_file.replace('.', '_bead_testBN_PSNR{:.2f}.'.format(psnr)))
